# coding: UTF-8
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import BertModel, BertTokenizer


class Model(nn.Module):

    def __init__(self, args):
        super(Model, self).__init__()
        self.bert = BertModel.from_pretrained(args.bert_path)
        for param in self.bert.parameters():
            param.requires_grad = True
        filter_num = args.filter_num
        filter_sizes  = [int(size) for size in args.filter_sizes.split(',')]
        self.convs = nn.ModuleList([
            nn.Conv2d(in_channels=1,
                      out_channels=filter_num,
                      kernel_size=(fs, args.hidden_size))
            for fs in filter_sizes
        ])
        self.dropout = nn.Dropout(args.dropout)

        self.fc_cnn = nn.Linear(len(filter_sizes) * filter_num, args.class_num)

    # def conv_and_pool(self, x, conv):
    #     x = F.relu(conv(x)).squeeze(3)
    #     x = F.max_pool1d(x, x.size(2)).squeeze(2)
    #     return x

    def forward(self, x):
        context = x[0]  # 输入的句子
        mask = x[2]  # 对padding部分进行mask，和句子一个size，padding部分用0表示，如：[1, 1, 1, 1, 0, 0]
        bertout= self.bert(context, attention_mask=mask)

        pooler_output= bertout['last_hidden_state']
        out = pooler_output.unsqueeze(1)
        # out = torch.cat([self.conv_and_pool(out, conv) for conv in self.convs], 1)

        # embedded = [batch size, 1, sent len, emb dim]
        # out1 = F.max_pool1d(torch.Size([32, 300, 220]),220).squeeze(2)

        conved = [F.relu(conv(out)).squeeze(3) for conv in self.convs]

        # conv_n = [batch size, n_filters, sent len - filter_sizes[n]]

        pooled = [F.max_pool1d(conv, conv.shape[2]).squeeze(2) for conv in conved]

        # pooled_n = [batch size, n_filters]

        out = self.dropout(torch.cat(pooled, dim=1))

        out = self.fc_cnn(out)
        return out